# ==============================================================================
# 文件: models.py
# 作用: 定义将作用于嵌入向量的分类器模型。
#      这里我们实现一个简单的多层感知机（MLP），它的结构可以灵活配置。
# ==============================================================================
import torch.nn as nn # 导入PyTorch的神经网络模块

# 定义一个MLP分类器类，它继承自PyTorch的nn.Module
class SimpleMLP(nn.Module):
    """
    一个简单的多层感知机（MLP）分类器。
    输入: 嵌入向量
    输出: 一个原始的、未经过sigmoid激活的logit值
    """
    def __init__(self, input_dim, hidden_dims):
        """
        构造函数，用于初始化MLP的结构。
        :param input_dim: int, 输入嵌入向量的维度。
        :param hidden_dims: list of int, 每个隐藏层的神经元数量。
        """
        super(SimpleMLP, self).__init__() # 调用父类的构造函数

        layers = [] # 创建一个列表来存放网络的所有层
        # 将输入维度加入到维度列表中，方便循环创建网络层
        all_dims = [input_dim] + hidden_dims

        # 循环创建隐藏层
        for i in range(len(all_dims) - 1):
            # 添加一个线性层（全连接层）
            layers.append(nn.Linear(all_dims[i], all_dims[i+1]))
            # 添加一个ReLU激活函数，增加网络的非线性
            layers.append(nn.ReLU())
            # 添加一个Dropout层，防止过拟合
            layers.append(nn.Dropout(0.3))

        # 添加最后的输出层，输出一个单独的logit值
        layers.append(nn.Linear(all_dims[-1], 1))

        # 使用nn.Sequential将所有层按顺序组合成一个网络模型
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        """
        定义模型的前向传播逻辑。
        :param x: Tensor, 输入的嵌入向量批次。
        :return: Tensor, 输出的logit值批次。
        """
        return self.network(x) # 将输入x通过网络，得到输出
